#include "trajectory_prediction.hpp"

int HitFlag = 0;

/**
 * @brief 公式拟合轨迹预测
 * @param 输入球心位置，时间戳
 * @return 拟合后的方程系数
 * @details 
*/
bool TraPredict::Ball_trajectory(const PointType center_point, const timeval& time, Eigen::VectorXf& result_x, Eigen::VectorXf& result_y, Eigen::VectorXf& result_z)
{
    // 所有非零点都放入容器中
    if(center_point.x != 0 && center_point.y != 0 && center_point.z != 0)
    {
        tra_point_cloud->points.push_back(center_point);
        tra_time_sec.push_back(time.tv_sec % 10000 + int(time.tv_usec / 1000)*0.001);
        // 把轨迹范围内的点单独保存
        if(detect_flag == 1)
        {
            int point_cloud_num = tra_point_cloud->size();
            time_sec.push_back(tra_time_sec.at(point_cloud_num - 1) - start_time);
            ball_tra_x.push_back(tra_point_cloud->at(point_cloud_num - 1).x);
            ball_tra_y.push_back(tra_point_cloud->at(point_cloud_num - 1).y);
            ball_tra_z.push_back(tra_point_cloud->at(point_cloud_num - 1).z);
        }
    }

    // 使用击球平台动作作为轨迹起点标志
    if(HitFlag != 0 && detect_flag == 0){
        start_time = 0;
        time_sec.clear();
        ball_tra_x.clear();
        ball_tra_y.clear();
        ball_tra_z.clear();
        result_x.setZero();
        result_y.setZero();
        result_z.setZero();
        tra_output.clear();

        detect_flag = 1;
        start_time = tra_time_sec.back();
    }
    if(HitFlag == 0 && detect_flag == 1 && tra_output.current_speed_y > 0 && tra_output.current_y > -0.7){
        detect_flag = 0;
    }

    // // 依据已有点进行轨迹起点和终点判断
    // if(start_flag(tra_point_cloud, tra_time_sec) > 0 && detect_flag == 0)
    // {
    //     int start_index = start_flag(tra_point_cloud, tra_time_sec);
    //     detect_flag = 1;
    //     start_time = tra_time_sec.at(start_index);
    //     // 保存起点到当前点的数据
    //     for(int it = start_index; it < tra_point_cloud->size(); it++)
    //     {
    //         time_sec.push_back(tra_time_sec.at(it) - start_time);
    //         ball_tra_x.push_back(tra_point_cloud->at(it).x);
    //         ball_tra_y.push_back(tra_point_cloud->at(it).y);
    //         ball_tra_z.push_back(tra_point_cloud->at(it).z);
    //     }
    //     cout << "start_flag:    " << start_index << endl;
    //     cout << tra_point_cloud->at(start_index) << endl;
    // }
    // if(end_flag(tra_point_cloud, tra_time_sec) > 0 && detect_flag == 1)
    // {
    //     int end_index = end_flag(tra_point_cloud, tra_time_sec);
    //     detect_flag = 0;
    //     start_time = 0;
    //     time_sec.clear();
    //     ball_tra_x.clear();
    //     ball_tra_y.clear();
    //     ball_tra_z.clear();
    //     result_x.setZero();
    //     result_y.setZero();
    //     result_z.setZero();
    //     tra_output.clear();
    //     cout << "end_flag:    " << end_index << endl;
    //     cout << tra_point_cloud->at(end_index) << endl;
    // }


    // 轨迹点大于8个开始预测
    if(ball_tra_x.size() >= PRED_THRES)
    {
        // // 清零检测标志，在使用击球平台动作作为轨迹起点标志时，需要在此处清零检测标志位
        // detect_flag = 0;
        // 裁剪轨迹点，只取最后n个
        vector<float> time_sec_clip;
        vector<float> ball_tra_x_clip;
        vector<float> ball_tra_y_clip;
        vector<float> ball_tra_z_clip;
        for(int i = time_sec.size() - CLIP_NUM; i < time_sec.size(); i++){
            time_sec_clip.push_back(time_sec[i]);
            ball_tra_x_clip.push_back(ball_tra_x[i]);
            ball_tra_y_clip.push_back(ball_tra_y[i]);
            ball_tra_z_clip.push_back(ball_tra_z[i]);
            cout << "time:  " << time_sec.at(i) << "   x_clip: " << ball_tra_x.at(i) << endl;
        }

        result_x = FitterLeastSquareMethod(time_sec_clip, ball_tra_x_clip, 1);
        result_y = FitterLeastSquareMethod(time_sec_clip, ball_tra_y_clip, 2);
        result_z = FitterLeastSquareMethod(time_sec_clip, ball_tra_z_clip, 1);

        // // 不做裁剪全部点进行预测
        // result_x = FitterLeastSquareMethod(time_sec, ball_tra_x, 1);
        // result_y = FitterLeastSquareMethod(time_sec, ball_tra_y, 2);
        // result_z = FitterLeastSquareMethod(time_sec, ball_tra_z, 1);

        cout << "\nThe points used: \n" << endl;
        for(int it = 0; it != ball_tra_x.size(); ++it)
        {
            cout << "time:  " << time_sec.at(it) << "   x: " << ball_tra_x.at(it) << endl;
        }
        // cout << "ball_tra_size:" << ball_tra_x.size() << endl;
        return true;
    }

    return false;
}

/**
 * @brief 根据轨迹方程计算所需的数据
 * @param 时间戳，三个轴的轨迹方程系数
 * @return 一个包含时间、速度，坐标，偏航角，俯仰角，最高点的结构体
 * @details 偏航角和俯仰角为角度制
*/
bool TraPredict::Output_process(const timeval& time, const Eigen::VectorXf& result_x, const Eigen::VectorXf& result_y, const Eigen::VectorXf& result_z, OutputValue& output)
{
    float end_time = 0;
    float current_time = 0;
    float max_time = 0;
    float predicted_speed_x = 0;
    float predicted_speed_y = 0;
    float predicted_speed_z = 0;

    if(tra_point_cloud->size() < 2) return false;
    output.current_x = tra_point_cloud->back().x;
    output.current_y = tra_point_cloud->back().y;
    output.current_z = tra_point_cloud->back().z;
    int it = tra_point_cloud->size() - 1;
    output.current_speed_x = (tra_point_cloud->at(it).x - tra_point_cloud->at(it - 1).x) / (tra_time_sec.at(it) - tra_time_sec.at(it - 1));
    output.current_speed_y = (tra_point_cloud->at(it).y - tra_point_cloud->at(it - 1).y) / (tra_time_sec.at(it) - tra_time_sec.at(it - 1));
    output.current_speed_z = (tra_point_cloud->at(it).z - tra_point_cloud->at(it - 1).z) / (tra_time_sec.at(it) - tra_time_sec.at(it - 1));
    float current_speed = pow(pow(output.current_speed_x, 2) + pow(output.current_speed_y, 2) + pow(output.current_speed_z, 2), 0.5);
    output.current_pitch = -acos(output.current_speed_y / current_speed) / M_PI * 180 + 90;
    output.current_yaw = atan2(output.current_speed_z, output.current_speed_x) / M_PI * 180 + 90;

    if(result_y[2] < 2)return false;
    float root_adjustment = pow(result_y[1], 2) - 4 * result_y[2] * (result_y[0] - DETECT_PLANE);
    if(root_adjustment < 0)return false;
    if(root_adjustment == 0)
    {
        end_time = (-0.5) * result_y[1] / result_y[2];
        max_time = end_time;
    }
    if(root_adjustment > 0)
    {
        end_time = (-result_y[1] + pow((pow(result_y[1], 2) - 4 * result_y[2] * (result_y[0] - DETECT_PLANE)), 0.5)) / (2 * result_y[2]);
        max_time = (-0.5) * result_y[1] / result_y[2];
    }
    cout << "end_time:  " << end_time << endl;
    current_time = (time.tv_sec % 10000 + int(time.tv_usec / 1000)*0.001) - start_time;
    output.delta_time = end_time - current_time;
    cout << "delta_time:  " << output.delta_time << endl;
    output.end_x = result_x[1] * end_time + result_x[0];
    output.end_z = result_z[1] * end_time + result_z[0];
    cout << "end_x_z:  " << output.end_x << "   " << output.end_z << endl;
    output.max_y = result_y[2] * max_time * max_time + result_y[1] * max_time + result_y[0];
    cout << "max_y:  " << output.max_y << endl;
    predicted_speed_x = result_x[1];
    predicted_speed_y = 2 * result_y[2] * end_time +result_y[1];
    predicted_speed_z = result_z[1];
    output.end_speed_x = predicted_speed_x;
    output.end_speed_y = predicted_speed_y;
    output.end_speed_z = predicted_speed_z;
    output.end_speed = pow(pow(predicted_speed_x, 2) + pow(predicted_speed_y, 2) + pow(predicted_speed_z, 2), 0.5);
    cout << "end_speed:  " << output.end_speed << endl;
    output.end_pitch = -acos(predicted_speed_y / output.end_speed) / M_PI * 180 + 90;
    output.end_yaw = atan2(predicted_speed_z, predicted_speed_x) / M_PI * 180 + 90;
    cout << "end_pitch:  " << output.end_pitch << endl;
    cout << "end_yaw:  " << output.end_yaw << endl;
    // output.current_speed_x = result_x[1];
    // output.current_speed_y = 2 * result_y[2] * current_time +result_y[1];
    // output.current_speed_z = result_z[1];
    output.result_x_1 = result_x[1];
    output.result_x_0 = result_x[0];
    output.result_y_2 = result_y[2];
    output.result_y_1 = result_y[1];
    output.result_y_0 = result_y[0];
    output.result_z_1 = result_z[1];
    output.result_z_0 = result_z[0];
}


/**
 * @brief 轨迹起点判断
 * @param 输入所有球心的点云和对应时间戳数据
 * @return 单个轨迹的起始点索引
 * @details 连续多个同方向速度则视为轨迹起始
*/
int TraPredict::start_flag(pPointCloud tra_point_cloud, vector<float> tra_time_sec)
{
    int start_index = 0;
    // 大于10个点开始计算速度，连续6个点Y轴速度方向一致，则输出起始点索引
    if(tra_point_cloud->size() > 6)
    {
        int detect_num = 6;
        int positive_num = 0;
        int negative_num = 0;
        Point3f temp_ball_speed(0,0,0);
        for(int it = tra_point_cloud->size() - 1; it > tra_point_cloud->size() - detect_num - 1; --it)
        {
            temp_ball_speed.x = (tra_point_cloud->at(it).x - tra_point_cloud->at(it - 1).x)
                                /(tra_time_sec.at(it) - tra_time_sec.at(it - 1));
            temp_ball_speed.y = (tra_point_cloud->at(it).y - tra_point_cloud->at(it - 1).y)
                                /(tra_time_sec.at(it) - tra_time_sec.at(it - 1));
            temp_ball_speed.z = (tra_point_cloud->at(it).z - tra_point_cloud->at(it - 1).z)
                                /(tra_time_sec.at(it) - tra_time_sec.at(it - 1));
            if(temp_ball_speed.y > 0 && (tra_time_sec.at(it) - tra_time_sec.at(it - 1)) < 0.1) positive_num ++;
            if(temp_ball_speed.y < 0 && (tra_time_sec.at(it) - tra_time_sec.at(it - 1)) < 0.1) negative_num ++;
        }
        if(positive_num == detect_num || negative_num == detect_num)
        {
            start_index = tra_point_cloud->size() - detect_num - 1;
        }
    }
    return start_index;
}

/**
 * @brief 轨迹终点判断
 * @param 输入所有球心的点云和对应时间戳数据
 * @return 单个轨迹的终点索引
 * @details 连续多个同方向速度后突然出现相反方向速度视为终点
*/
int TraPredict::end_flag(pPointCloud tra_point_cloud, vector<float> tra_time_sec)
{
    int end_index = 0;
    // 大于10个点开始计算速度，连续3个点Y轴速度方向一致，再连续3个点速度方向相反，则输出终点索引
    if(tra_point_cloud->size() > 10)
    {
        int detect_num = 6;
        int half_num = 3;
        int positive_num1 = 0;
        int negative_num1 = 0;
        int positive_num2 = 0;
        int negative_num2 = 0;
        Point3f temp_ball_speed(0,0,0);
        for(int it = tra_point_cloud->size() - 1; it > tra_point_cloud->size() - half_num - 1; --it)
        {
            temp_ball_speed.x = (tra_point_cloud->at(it).x - tra_point_cloud->at(it - 1).x)
                                /(tra_time_sec.at(it) - tra_time_sec.at(it - 1));
            temp_ball_speed.y = (tra_point_cloud->at(it).y - tra_point_cloud->at(it - 1).y)
                                /(tra_time_sec.at(it) - tra_time_sec.at(it - 1));
            temp_ball_speed.z = (tra_point_cloud->at(it).z - tra_point_cloud->at(it - 1).z)
                                /(tra_time_sec.at(it) - tra_time_sec.at(it - 1));
            if(temp_ball_speed.y > 0)positive_num1 ++;
            if(temp_ball_speed.y < 0)negative_num1 ++;
        }
        for(int it = tra_point_cloud->size() - half_num - 1; it > tra_point_cloud->size() - detect_num - 1; --it)
        {
            temp_ball_speed.x = (tra_point_cloud->at(it).x - tra_point_cloud->at(it - 1).x)
                                /(tra_time_sec.at(it) - tra_time_sec.at(it - 1));
            temp_ball_speed.y = (tra_point_cloud->at(it).y - tra_point_cloud->at(it - 1).y)
                                /(tra_time_sec.at(it) - tra_time_sec.at(it - 1));
            temp_ball_speed.z = (tra_point_cloud->at(it).z - tra_point_cloud->at(it - 1).z)
                                /(tra_time_sec.at(it) - tra_time_sec.at(it - 1));
            if(temp_ball_speed.y > 0)positive_num2 ++;
            if(temp_ball_speed.y < 0)negative_num2 ++;
        }
        if(positive_num2 == half_num && negative_num1 == half_num)//((positive_num1 == half_num && negative_num2 == half_num) || (positive_num2 == half_num && negative_num1 == half_num))
        {
            end_index = tra_point_cloud->size() - half_num - 1;
        }
        // 若长时间未刷新点，则输出终点
        if(tra_time_sec.at(tra_time_sec.size() - 1) - tra_time_sec.at(tra_time_sec.size() - 2) > 0.1)
        {
            end_index = tra_point_cloud->size() - 2;
        }
    }
    return end_index;
}

/**
 * @brief 最小二乘法拟合多项式
 * @param X X轴数据 
 * @param Y Y轴数据
 * @param orders 拟合方程的阶次，需大于零
 * @return Eigen::VectorXf 拟合多项式的系数向量
 */
Eigen::VectorXf TraPredict::FitterLeastSquareMethod(vector<float> &X, vector<float> &Y, uint8_t orders)
{
    // abnormal input verification
    if (X.size() < 2 || Y.size() < 2 || X.size() != Y.size() || orders < 1)
        exit(EXIT_FAILURE);

    // map sample data from STL vector to eigen vector
    Eigen::Map<Eigen::VectorXf> sampleX(X.data(), X.size());
    Eigen::Map<Eigen::VectorXf> sampleY(Y.data(), Y.size());

    Eigen::MatrixXf mtxVandermonde(X.size(), orders + 1);  // Vandermonde matrix of X-axis coordinate vector of sample data
    Eigen::VectorXf colVandermonde = sampleX;              // Vandermonde column

    // construct Vandermonde matrix column by column
    for (size_t i = 0; i < orders + 1; ++i)
    {
        if (0 == i)
        {
            mtxVandermonde.col(0) = Eigen::VectorXf::Constant(X.size(), 1, 1);
            continue;
        }
        if (1 == i)
        {
            mtxVandermonde.col(1) = colVandermonde;
            continue;
        }
        colVandermonde = colVandermonde.array()*sampleX.array();
        mtxVandermonde.col(i) = colVandermonde;
    }

    // calculate coefficients vector of fitted polynomial
    Eigen::VectorXf result = (mtxVandermonde.transpose()*mtxVandermonde).inverse()*(mtxVandermonde.transpose())*sampleY;

    return result;
}

TraPredict::TraPredict():
tra_point_cloud(new PointCloud),
detect_flag(0),
result_x(2),
result_y(3),
result_z(2)
{
    result_x << 0, 0;
    result_y << 0, 0, 0;
    result_z << 0, 0;
}
TraPredict::~TraPredict()
{

}